#include "codegen.h"
#include "llvm/IR/Verifier.h"

/*
; ModuleID = 'expr'
source_filename = "expr"

@0 = private unnamed_addr constant [16 x i8] c"expr value: %d\0A\00", align 1

declare i32 @printf(ptr, ...)

define i32 @main() {
entry:
  %a = alloca i32, align 4
  store i32 0, ptr %a, align 4
  %a1 = load i32, ptr %a, align 4
  %b = alloca i32, align 4
  store i32 5, ptr %b, align 4
  %b2 = load i32, ptr %b, align 4
  br label %cond

cond:                                             ; preds = %entry
  %a3 = load i32, ptr %a, align 4
  %0 = icmp ne i32 %a3, 0
  br i1 %0, label %then, label %else

then:                                             ; preds = %cond
  store i32 3, ptr %b, align 4
  %b4 = load i32, ptr %b, align 4
  br label %last

else:                                             ; preds = %cond
  store i32 4, ptr %a, align 4
  %a5 = load i32, ptr %a, align 4
  br label %last

last:                                             ; preds = %else, %then
  %a6 = load i32, ptr %a, align 4
  %b7 = load i32, ptr %b, align 4
  %mul = mul nsw i32 %a6, %b7
  %sub = sub nsw i32 %mul, 4
  %1 = call i32 (ptr, ...) @printf(ptr @0, i32 %sub)
  ret i32 0
}

*/

llvm::Value *CodeGen::VisitProgram(Program *p)
{
    // 创建 printf 函数
    auto printFunctionType = llvm::FunctionType::get(irBuilder.getInt32Ty(), {irBuilder.getInt8PtrTy()}, true);
    auto printFunction = llvm::Function::Create(printFunctionType, llvm::GlobalValue::ExternalLinkage, "printf", module.get());
    // 创建 main 函数
    auto mainFunctionType = llvm::FunctionType::get(irBuilder.getInt32Ty(), false);
    auto mainFunction = llvm::Function::Create(mainFunctionType, llvm::GlobalValue::ExternalLinkage, "main", module.get());
    // 创建 main 函数的基本块
    llvm::BasicBlock *entryBasicBlock = llvm::BasicBlock::Create(context, "entry", mainFunction);
    // 设置该基本块作为指令的入口
    irBuilder.SetInsertPoint(entryBasicBlock);
    // 记录当前函数
    curFunc = mainFunction;

    llvm::Value *lastVal = nullptr;
    for (auto node : p->nodeVec)
    {
        lastVal = node->Accept(this);
    }
    if (lastVal)
        irBuilder.CreateCall(printFunction, {irBuilder.CreateGlobalStringPtr("expr value: %d\n"), lastVal});
    else
        irBuilder.CreateCall(printFunction, {irBuilder.CreateGlobalStringPtr("last instruction is not expr!\n")});

    // 创建返回值
    llvm::Value *ret = irBuilder.CreateRet(irBuilder.getInt32(0));

    llvm::verifyFunction(*mainFunction);

    module->print(llvm::outs(), nullptr);
    return ret;
}

llvm::Value *CodeGen::VisitDeclStmt(DeclStmt *declStmt)
{
    llvm::Value *lastVal = nullptr;
    for (auto node : declStmt->nodeVec)
    {
        lastVal = node->Accept(this);
    }
    return lastVal;
}

llvm::Value *CodeGen::VisitBlockStmt(BlockStmt *block)
{
    llvm::Value *lastVal = nullptr;
    for (auto node : block->nodeVec)
    {
        lastVal = node->Accept(this);
    }
    return lastVal;
}

llvm::Value *CodeGen::VisitIfStmt(IfStmt *ifStmt)
{
    llvm::BasicBlock *condBasicBlock = llvm::BasicBlock::Create(context, "cond", curFunc);
    llvm::BasicBlock *thenBasicBlock = llvm::BasicBlock::Create(context, "then", curFunc);
    llvm::BasicBlock *elseBasicBlock = nullptr;
    if (ifStmt->elseNode)
        elseBasicBlock = llvm::BasicBlock::Create(context, "else", curFunc);
    llvm::BasicBlock *lastBasicBlock = llvm::BasicBlock::Create(context, "last", curFunc);

    // 需要手动添加一个无条件跳转指令，llvm 不会自动完成这个工作
    irBuilder.CreateBr(condBasicBlock);
    irBuilder.SetInsertPoint(condBasicBlock);
    llvm::Value *ret = ifStmt->condNode->Accept(this);
    // 整形比较指令
    llvm::Value *condVal = irBuilder.CreateICmpNE(ret, irBuilder.getInt32(0)); // 这里需要判断条件是否为真

    if (ifStmt->elseNode)
    {
        irBuilder.CreateCondBr(condVal, thenBasicBlock, elseBasicBlock);

        // handle then basic block
        irBuilder.SetInsertPoint(thenBasicBlock);
        ifStmt->thenNode->Accept(this);
        irBuilder.CreateBr(lastBasicBlock);

        // handle else basic block
        irBuilder.SetInsertPoint(elseBasicBlock);
        ifStmt->elseNode->Accept(this);
        irBuilder.CreateBr(lastBasicBlock);
    }
    else
    {
        irBuilder.CreateCondBr(condVal, thenBasicBlock, lastBasicBlock);

        // handle then basic block
        irBuilder.SetInsertPoint(thenBasicBlock);
        ifStmt->thenNode->Accept(this);
        irBuilder.CreateBr(lastBasicBlock);
    }
    irBuilder.SetInsertPoint(lastBasicBlock);
    return nullptr;
}

llvm::Value *CodeGen::VisitVariableDecl(VariableDecl *decl)
{
    llvm::Type *ty = nullptr;
    if (decl->type == CType::getIntTy())
    {
        ty = irBuilder.getInt32Ty();
    }
    llvm::StringRef text(decl->token.ptr, decl->token.length);
    llvm::Value *varAddr = irBuilder.CreateAlloca(ty, nullptr, text);
    varAddrTyMap.insert({text, {varAddr, ty}});
    return varAddr;
}

llvm::Value *CodeGen::VisitVariableAccessExpr(VariableAccessExpr *varaccExpr)
{
    llvm::StringRef text(varaccExpr->token.ptr, varaccExpr->token.length);
    std::pair pair = varAddrTyMap[text];
    llvm::Value *varAddr = pair.first;
    llvm::Type *ty = pair.second;
    if (varaccExpr->type == CType::getIntTy())
    {
        ty = irBuilder.getInt32Ty();
    }
    // 返回一个右值
    return irBuilder.CreateLoad(ty, varAddr, text);
}

// a = 3; // right value
llvm::Value *CodeGen::VisitAssignExpr(AssignExpr *assignExpr)
{
    VariableAccessExpr *varAccExpr = (VariableAccessExpr *)assignExpr->left.get();
    llvm::StringRef text(varAccExpr->token.ptr, varAccExpr->token.length);
    std::pair pair = varAddrTyMap[text];
    llvm::Value *addr = pair.first;
    llvm::Type *ty = pair.second;
    llvm::Value *rValue = assignExpr->right->Accept(this);
    // 这个得到的是一个左值
    irBuilder.CreateStore(rValue, addr);
    // 返回一个右值
    return irBuilder.CreateLoad(ty, addr, text);
}

llvm::Value *CodeGen::VisitBinaryExpr(BinaryExpr *binaryExpr)
{
    auto left = binaryExpr->left->Accept(this);
    auto right = binaryExpr->right->Accept(this);

    switch (binaryExpr->op)
    {
    case OpCode::add:
    {
        return irBuilder.CreateNSWAdd(left, right, "add"); // CreateNSW... 是防止溢出行为的
    }
    case OpCode::sub:
    {
        return irBuilder.CreateNSWSub(left, right, "sub");
    }
    case OpCode::mul:
    {
        return irBuilder.CreateNSWMul(left, right, "mul");
    }
    case OpCode::div:
    {
        return irBuilder.CreateSDiv(left, right, "div");
    }
    default:
    {
        break;
    }
    }
    return nullptr;
}

llvm::Value *CodeGen::VisitNumberExpr(NumberExpr *factorExpr)
{
    return irBuilder.getInt32(factorExpr->token.value);
}